import os
from collections import defaultdict
if __name__ == '__main__':
    compression_time_list = defaultdict(list)
    decompression_time_list = defaultdict(list)
    is_add_to = 0
    enable_random = 0
    maximum_x_plus1 = 28
    for bitwidth in [2,4,8]:
        for x in range(5):
            compression_time_list[bitwidth].append(0)
            decompression_time_list[bitwidth].append(0)
        for x in range(5, maximum_x_plus1):
            os.system(f"./benchmark_terngrad_nvcc.o {x} {bitwidth} {is_add_to} {enable_random}")
            with open("ret.txt", 'r') as f:
                compression_time_list[bitwidth].append(float(f.readline().strip()))
                decompression_time_list[bitwidth].append(float(f.readline().strip()))
    
    with open("benchmark_terngrad.csv", "w") as f:
        f.write(f"""Benchmark Terngrad:
params: 
    is_add_to = {is_add_to}
    enable_random = {enable_random}
""")
        f.write(','.join(list(map(str,['CT(us)',2,4,8])))+'\n')
        for x in range(5, maximum_x_plus1):
            f.write(f"{x},")
            f.write(','.join(list(map(str,[compression_time_list[bitwidth][x] for bitwidth in [2,4,8]]))))
            f.write('\n')
        f.write('\n')
        f.write(','.join(list(map(str,['DT(us)',2,4,8])))+'\n')
        for x in range(5, maximum_x_plus1):
            f.write(f"{x},")
            f.write(','.join(list(map(str,[decompression_time_list[bitwidth][x] for bitwidth in [2,4,8]]))))
            f.write('\n')



            
            

    print(f"compression_time={compression_time_list}")
    print(f"decompression_time={decompression_time_list}")
        